import json
import datetime
import numpy as np
from scipy.interpolate import UnivariateSpline
import matplotlib.pyplot as plt
import argparse

def time2sec(t):
    return t.hour*3600 + t.minute*60 + t.second

def sec2time(s):
    hour = s//3600
    minute = (s%3600)//60
    second = s%60
    return datetime.time(hour,minute,second)
    
def calcu_density(chat_counts, window_size=10):
    time_lim = len(chat_counts)
    time_x = []
    density_y = []
    for i in range(0,time_lim+1,window_size):
        left_bound = max(0,int(i-window_size/2))
        right_bound = min(time_lim,int(i+window_size/2)+1)
        real_window_size = right_bound - left_bound
        time_x.append(i)
        density_y.append(sum(chat_counts[left_bound:right_bound])/real_window_size)
    return time_x, density_y

if __name__ == '__main__':
    parser = argparse.ArgumentParser("根据json文件画出chat的密度图。\n")
    parser.add_argument('path', type=str, help='chat的json文件路径')
    parser.add_argument('--smooth', type=int, default=50, help='平滑系数，建议在1-1000之间')
    parser.add_argument('--width', type=int, default=1920, help='smoothed图片的宽，单位:pixel')
    parser.add_argument('--height', type=int, default=100, help='smoothed图片的高，单位:pixel')
    parser.add_argument('--alpha', type=float, default=0.5, help='图片透明度')
    args = parser.parse_args()

    # args
    chat_file_path = args.path
    pic_width = args.width
    pic_height = args.height
    pic_dpi = 100 # pixels per inch
    smooth_factor = args.smooth # 1-1000 maybe
    pic_size = (pic_width/pic_dpi,pic_height/pic_dpi) # inch
    pic_alpha = args.alpha

    # 读文件
    with open(chat_file_path,'r',encoding='utf-8') as f:
        text = f.read()
        chat_list = json.loads(text)
    
    chat_secs = [time2sec(datetime.time.fromisoformat(chat['dt'])) for chat in chat_list]
    time_lim = chat_secs[-1]
    
    # 统计每秒的chat数量
    chat_counts = [0]*(time_lim+1)
    for chat_sec in chat_secs:
        chat_counts[chat_sec] += 1
    
    # 计算密度，平滑
    time_x, density_y = calcu_density(chat_counts)
    density_y = np.array(density_y)
    s = UnivariateSpline(time_x,density_y,s=smooth_factor)
    # cliped_density = density_y[density_y>2]
    # max_k_time = np.argsort(cliped_density)[:10]
    # max_k_time = [sec2time(time).isoformat() for time in max_k_time]

    # 画图
    fig = plt.figure(figsize=pic_size,dpi=pic_dpi)
    ax=plt.Axes(fig,[0., 0., 1., 1.])
    fig.add_axes(ax)
    plt.fill_between(time_x,0,density_y,where=density_y>=0,color='black',interpolate=True,alpha=pic_alpha)
    plt.xlim((0,time_x[-1]))
    plt.savefig('raw_pic.png', transparent=True, bbox_inches='tight', pad_inches=0.0)
    fig = plt.figure(figsize=pic_size,dpi=pic_dpi)
    ax=plt.Axes(fig,[0., 0., 1., 1.])
    ax.set_axis_off() 
    fig.add_axes(ax)
    plt.fill_between(time_x,0,s(time_x),where=s(time_x)>=0,color='black',interpolate=True,alpha=pic_alpha)
    plt.xlim((0,time_x[-1]))
    plt.savefig('smoothed_pic.png', transparent=True, bbox_inches='tight', pad_inches=0.0)

